from one_layer_model import IlluNet_with_Quad
from LOL import LolTrainLoader
from lr_scheduler import CosineAnnealingLR_Restart
class L1_loss(nn.Module):
    
    def __init__(self):
        super(L1_loss, self).__init__()
        self.l1_loss = nn.L1Loss()

    def forward(self, img, target):
        return self.l1_loss(img, target)
def train():
  
    gt_dir = 'xxx'
    img_dir = 'xxx'
    batch_size = 16
    epochs = 300
    device = torch.device('cuda') 
    criterion = L1_loss()
    criterion.train()
    model = IlluNet_with_Quad(3,3,None).to(device)
    
    model.train()
    optimizer = AdamW(model.parameters(), lr=2e-4,  betas=(0.9, 0.999))
    T_period = [40000, 40000, 40000,40000]
    restarts= [40000, 80000, 120000]
    restart_weights = [1, 1, 1]
    scheduler = CosineAnnealingLR_Restart( optimizer,T_period, eta_min=1e-6, restarts=restarts,
                                            weights=restart_weights)
    TrainDataset = LolTrainLoader( gt_dir,img_dir,patch_size)
    train_queue = torch.utils.data.DataLoader(
        TrainDataset,
        batch_size=batch_size,
        pin_memory=True,
        num_workers=2,
        shuffle=True,)

    total_step = 0
    for epoch in range(epochs):
        if total_step > T_period[0]:
            break
        model.train()
        losses = []
        
        for batch_idx, (dark_img,gt_img,name) in enumerate(train_queue):
            if total_step > T_period[0]-10:
                break
            total_step += 1
            gt_img = gt_img.to(device)
            dark_img = dark_img.to(device)
            optimizer.zero_grad()
            output = model(dark_img)

            loss_l1 = criterion_illu(output, gt_img)
            loss_l1.backward()
           
            optimizer.step()
            scheduler.step()
            
            base_lr = optimizer.param_groups[0]['lr']
            losses.append(loss_l1.item())
          
        
        torch.save(model, os.path.join(model_path, 'weights_last.pt'))
        if np.mean(losses)<best_loss:
         
            best_loss=np.mean(losses)
            torch.save(model, os.path.join(model_path, 'weights_best.pt'))

                
if __name__=='__main__':
  
    train()